# 人工智能NLP-Agent数字人项目-04-基金数据问答任务工单V1.1-2.13
import csv
import sqlite3
import pandas as pd

import utils.configFinRAG as configFinRAG


def query_db(sql, cursor):
    try:
        cursor.execute(sql)
        result_cols = cursor.fetchall()
        exc_result = str(result_cols)
        success_flag = 1
    except Exception as e:
        exc_result = f"error: {str(e)}"
        success_flag = 0
        print(f"Error executing SQL: {sql}\n{e}")
    return success_flag, exc_result


if __name__ == '__main__':
    # 打开生成的sql文件
    question_sql_df = pd.read_csv(configFinRAG.question_sql_path, delimiter=",", header=0)

    # 数据库路径
    db_path = '/Users/wanghr/Documents/八维研修/项目/999-项目/fay数字人/金融场景智能问答系统/bs_challenge_financial_14b_dataset/dataset/博金杯比赛数据.db'

    # 使用上下文管理器管理文件和数据库连接
    with open(configFinRAG.question_sql_check_path, 'w', newline='', encoding='utf-8-sig') as file:
        csvwriter = csv.writer(file)
        csvwriter.writerow(['问题id', '问题', 'SQL', 'flag', '执行结果'])

        with sqlite3.connect(db_path) as conn:
            cs = conn.cursor()

            # 执行SQL 并返回结果写入sql执行结果文件
            for _, row in question_sql_df.iterrows():
                if row['SQL'] != 'error':
                    success_flag, exc_result = query_db(row['SQL'], cs)
                    csvwriter.writerow([row['问题id'], row['问题'], row['SQL'], success_flag, exc_result])



